perf: Optimize CUTLASS MoE helper kernels for small-batch decode workloads by bkryu · Pull Request #3014 · flashinfer-ai/flashinfer

bkryu · 2026-04-08T04:49:43Z

📌 Description

Summary

Ported optimizations from TRTLLM and expanded a bit:

Remove unnecessary N-dim SF padding loops from expandInputRowsKernel and doActivationKernel — the CUTLASS grouped GEMM never reads scale factors beyond tokens_to_expert
Add M=0 early exit in computeStridesTmaWarpSpecializedKernel to skip stride/pointer setup for experts with no assigned tokens
Spread computeStrides work across multiple SMs (4 blocks instead of 1)
Add __launch_bounds__ to doActivationKernel for better register allocation
Add new test coverage for N-dim padding safety to confirm that skipping N-dim SF padding is safe.

Motivation

When running CUTLASS MoE during the generation phase, three MoE helper kernels were severely underutilized:

Kernel	Issue
`expandInputRowsKernel`	87% of time writing zero SF padding for TMA alignment
`doActivationKernel`	63% of time writing zero SF padding
`computeStridesTmaWarpSpecializedKernel`	grid=1 (single SM), no early exit for 120 empty experts

The N-dim SF padding loops iterated over MinNDimAlignment x num_experts = 128 x 128 = 16384 potential padding slots to zero scale factors for token rows beyond each expert's actual token count. This padding was unnecessary — the CUTLASS grouped GEMM sets gemm_m = tokens_to_expert per expert and never reads scale factors for rows beyond that boundary.

Changes

expandInputRowsKernel + doActivationKernel -- Remove N-dim SF padding:
Deleted the entire N-dim SF padding section (after griddepcontrol.launch_dependents) from both kernels. The CUTLASS grouped GEMM's problem shapes bound the MMA tile access to tokens_to_expert rows per expert; padding rows are never read. Removed dead num_padding_tokens variables from both launchers and simplified grid formulas to be driven purely by the expanded token count.

K-dim SF padding (inside the per-token main loop) is preserved — MMA tiles can straddle the inter_size boundary within valid rows, requiring those positions to be zeroed.

doActivationKernel -- __launch_bounds__:
Added __launch_bounds__(ACTIVATION_THREADS_PER_BLOCK) to help the compiler optimize register allocation.

computeStridesTmaWarpSpecializedKernel -- M=0 early exit:
After writing problem_shapes[expert] and int4_groupwise_params.shape (which CUTLASS needs for all experts to traverse the problem list), experts with gemm_m == 0 return immediately — skipping setupFP4BlockScalingFactors, computeTmaWarpSpecializedInputStrides, and computeTmaWarpSpecializedInputPointers for both GEMMs. For decode with 128 experts and top_k=8, this skips ~120 experts' full setup.

computeStridesTmaWarpSpecializedKernel -- block size:
Changed std::min(1024, num_experts_per_node) to std::min(32, num_experts_per_node), spreading 128 experts across 4 blocks on 4 SMs instead of 1 block on 1 SM.

Correctness

The N-dim padding removal was validated with 0xFF poisoning: the SF buffer was filled with 0xFF (worst-case FP8 scale factor values of +/-448) before the kernel wrote real SFs to valid positions. With padding positions containing 0xFF, all tests pass with the same error rates as the original code — confirming the CUTLASS GEMM never reads the padding positions.

K-dim SF padding is preserved because MMA tiles straddle the K boundary within valid rows. The original author confirmed: "Tests should fail if the K dimension padding is disabled, but not if the N dimension stuff is."

All changes preserve PDL (griddepcontrol) overlap — no cudaMemsetAsync or stream operations are introduced.

Large-batch regression safety

All optimizations are unconditionally beneficial or neutral for large batches:

N-dim padding removal: unconditionally removes dead code — the padding was unnecessary at all batch sizes
Grid formulas: at large batch, expanded_tokens dominates and the grid remains at smCount * 8
M=0 early exit: doesn't trigger when experts have many tokens
__launch_bounds__: unconditionally better

Performance Numbers

Click to view `flashinfer_benchmark.py` test cases used to collect the data

--routine cutlass_fused_moe --num_tokens 1 --hidden_size 5120 --intermediate_size 8192 --num_experts 128 --top_k 1 --cutlass_variant base --input_dtype float16 --use_cuda_events --no_cuda_graph
--routine cutlass_fused_moe --num_tokens 1 --hidden_size 5120 --intermediate_size 8192 --num_experts 128 --top_k 1 --cutlass_variant fp8 --input_dtype float16 --use_cuda_events --no_cuda_graph
--routine cutlass_fused_moe --num_tokens 1 --hidden_size 5120 --intermediate_size 8192 --num_experts 128 --top_k 1 --cutlass_variant nvfp4 --input_dtype float16 --use_cuda_events --no_cuda_graph
--routine cutlass_fused_moe --num_tokens 1 --hidden_size 5120 --intermediate_size 8192 --num_experts 128 --top_k 1 --cutlass_variant nvfp4 --quantized_input --input_dtype float16 --use_cuda_events --no_cuda_graph
--routine cutlass_fused_moe --num_tokens 1 --hidden_size 5120 --intermediate_size 8192 --num_experts 128 --top_k 1 --cutlass_variant base --input_dtype float16 --ep_size 4 --ep_rank 0 --use_cuda_events --no_cuda_graph
--routine cutlass_fused_moe --num_tokens 1 --hidden_size 5120 --intermediate_size 8192 --num_experts 128 --top_k 1 --cutlass_variant base --input_dtype float16 --tp_size 4 --tp_rank 0 --use_cuda_events --no_cuda_graph
--routine cutlass_fused_moe --num_tokens 128 --hidden_size 5120 --intermediate_size 8192 --num_experts 128 --top_k 1 --cutlass_variant base --input_dtype float16 --use_cuda_events --no_cuda_graph
--routine cutlass_fused_moe --num_tokens 128 --hidden_size 5120 --intermediate_size 8192 --num_experts 128 --top_k 1 --cutlass_variant fp8 --input_dtype float16 --use_cuda_events --no_cuda_graph
--routine cutlass_fused_moe --num_tokens 128 --hidden_size 5120 --intermediate_size 8192 --num_experts 128 --top_k 1 --cutlass_variant nvfp4 --input_dtype float16 --use_cuda_events --no_cuda_graph
--routine cutlass_fused_moe --num_tokens 128 --hidden_size 5120 --intermediate_size 8192 --num_experts 128 --top_k 1 --cutlass_variant nvfp4 --quantized_input --input_dtype float16 --use_cuda_events --no_cuda_graph
--routine cutlass_fused_moe --num_tokens 128 --hidden_size 5120 --intermediate_size 8192 --num_experts 128 --top_k 1 --cutlass_variant base --input_dtype float16 --ep_size 4 --ep_rank 0 --use_cuda_events --no_cuda_graph
--routine cutlass_fused_moe --num_tokens 128 --hidden_size 5120 --intermediate_size 8192 --num_experts 128 --top_k 1 --cutlass_variant base --input_dtype float16 --tp_size 4 --tp_rank 0 --use_cuda_events --no_cuda_graph
--routine cutlass_fused_moe --num_tokens 1024 --hidden_size 5120 --intermediate_size 8192 --num_experts 128 --top_k 1 --cutlass_variant base --input_dtype float16 --use_cuda_events --no_cuda_graph
--routine cutlass_fused_moe --num_tokens 1024 --hidden_size 5120 --intermediate_size 8192 --num_experts 128 --top_k 1 --cutlass_variant fp8 --input_dtype float16 --use_cuda_events --no_cuda_graph
--routine cutlass_fused_moe --num_tokens 1024 --hidden_size 5120 --intermediate_size 8192 --num_experts 128 --top_k 1 --cutlass_variant nvfp4 --input_dtype float16 --use_cuda_events --no_cuda_graph
--routine cutlass_fused_moe --num_tokens 1024 --hidden_size 5120 --intermediate_size 8192 --num_experts 128 --top_k 1 --cutlass_variant nvfp4 --quantized_input --input_dtype float16 --use_cuda_events --no_cuda_graph
--routine cutlass_fused_moe --num_tokens 1024 --hidden_size 5120 --intermediate_size 8192 --num_experts 128 --top_k 1 --cutlass_variant base --input_dtype float16 --ep_size 4 --ep_rank 0 --use_cuda_events --no_cuda_graph
--routine cutlass_fused_moe --num_tokens 1024 --hidden_size 5120 --intermediate_size 8192 --num_experts 128 --top_k 1 --cutlass_variant base --input_dtype float16 --tp_size 4 --tp_rank 0 --use_cuda_events --no_cuda_graph
--routine cutlass_fused_moe --num_tokens 1 --hidden_size 7168 --intermediate_size 2048 --num_experts 256 --top_k 8 --cutlass_variant base --input_dtype float16 --use_cuda_events --no_cuda_graph
--routine cutlass_fused_moe --num_tokens 1 --hidden_size 7168 --intermediate_size 2048 --num_experts 256 --top_k 8 --cutlass_variant fp8 --input_dtype float16 --use_cuda_events --no_cuda_graph
--routine cutlass_fused_moe --num_tokens 1 --hidden_size 7168 --intermediate_size 2048 --num_experts 256 --top_k 8 --cutlass_variant nvfp4 --input_dtype float16 --use_cuda_events --no_cuda_graph
--routine cutlass_fused_moe --num_tokens 1 --hidden_size 7168 --intermediate_size 2048 --num_experts 256 --top_k 8 --cutlass_variant nvfp4 --quantized_input --input_dtype float16 --use_cuda_events --no_cuda_graph
--routine cutlass_fused_moe --num_tokens 1 --hidden_size 7168 --intermediate_size 2048 --num_experts 256 --top_k 8 --cutlass_variant base --input_dtype float16 --ep_size 4 --ep_rank 0 --use_cuda_events --no_cuda_graph
--routine cutlass_fused_moe --num_tokens 1 --hidden_size 7168 --intermediate_size 2048 --num_experts 256 --top_k 8 --cutlass_variant base --input_dtype float16 --tp_size 4 --tp_rank 0 --use_cuda_events --no_cuda_graph
--routine cutlass_fused_moe --num_tokens 128 --hidden_size 7168 --intermediate_size 2048 --num_experts 256 --top_k 8 --cutlass_variant base --input_dtype float16 --use_cuda_events --no_cuda_graph
--routine cutlass_fused_moe --num_tokens 128 --hidden_size 7168 --intermediate_size 2048 --num_experts 256 --top_k 8 --cutlass_variant fp8 --input_dtype float16 --use_cuda_events --no_cuda_graph
--routine cutlass_fused_moe --num_tokens 128 --hidden_size 7168 --intermediate_size 2048 --num_experts 256 --top_k 8 --cutlass_variant nvfp4 --input_dtype float16 --use_cuda_events --no_cuda_graph
--routine cutlass_fused_moe --num_tokens 128 --hidden_size 7168 --intermediate_size 2048 --num_experts 256 --top_k 8 --cutlass_variant nvfp4 --quantized_input --input_dtype float16 --use_cuda_events --no_cuda_graph
--routine cutlass_fused_moe --num_tokens 128 --hidden_size 7168 --intermediate_size 2048 --num_experts 256 --top_k 8 --cutlass_variant base --input_dtype float16 --ep_size 4 --ep_rank 0 --use_cuda_events --no_cuda_graph
--routine cutlass_fused_moe --num_tokens 128 --hidden_size 7168 --intermediate_size 2048 --num_experts 256 --top_k 8 --cutlass_variant base --input_dtype float16 --tp_size 4 --tp_rank 0 --use_cuda_events --no_cuda_graph
--routine cutlass_fused_moe --num_tokens 1024 --hidden_size 7168 --intermediate_size 2048 --num_experts 256 --top_k 8 --cutlass_variant base --input_dtype float16 --use_cuda_events --no_cuda_graph
--routine cutlass_fused_moe --num_tokens 1024 --hidden_size 7168 --intermediate_size 2048 --num_experts 256 --top_k 8 --cutlass_variant fp8 --input_dtype float16 --use_cuda_events --no_cuda_graph
--routine cutlass_fused_moe --num_tokens 1024 --hidden_size 7168 --intermediate_size 2048 --num_experts 256 --top_k 8 --cutlass_variant nvfp4 --input_dtype float16 --use_cuda_events --no_cuda_graph
--routine cutlass_fused_moe --num_tokens 1024 --hidden_size 7168 --intermediate_size 2048 --num_experts 256 --top_k 8 --cutlass_variant nvfp4 --quantized_input --input_dtype float16 --use_cuda_events --no_cuda_graph
--routine cutlass_fused_moe --num_tokens 1024 --hidden_size 7168 --intermediate_size 2048 --num_experts 256 --top_k 8 --cutlass_variant base --input_dtype float16 --ep_size 4 --ep_rank 0 --use_cuda_events --no_cuda_graph
--routine cutlass_fused_moe --num_tokens 1024 --hidden_size 7168 --intermediate_size 2048 --num_experts 256 --top_k 8 --cutlass_variant base --input_dtype float16 --tp_size 4 --tp_rank 0 --use_cuda_events --no_cuda_graph
--routine cutlass_fused_moe --num_tokens 1 --hidden_size 5120 --intermediate_size 8192 --num_experts 128 --top_k 1 --cutlass_variant base --input_dtype float16 --autotune --use_cuda_events --no_cuda_graph
--routine cutlass_fused_moe --num_tokens 1 --hidden_size 5120 --intermediate_size 8192 --num_experts 128 --top_k 1 --cutlass_variant fp8 --input_dtype float16 --autotune --use_cuda_events --no_cuda_graph
--routine cutlass_fused_moe --num_tokens 1 --hidden_size 5120 --intermediate_size 8192 --num_experts 128 --top_k 1 --cutlass_variant nvfp4 --input_dtype float16 --autotune --use_cuda_events --no_cuda_graph
--routine cutlass_fused_moe --num_tokens 1 --hidden_size 5120 --intermediate_size 8192 --num_experts 128 --top_k 1 --cutlass_variant nvfp4 --quantized_input --input_dtype float16 --autotune --use_cuda_events --no_cuda_graph
--routine cutlass_fused_moe --num_tokens 1 --hidden_size 5120 --intermediate_size 8192 --num_experts 128 --top_k 1 --cutlass_variant base --input_dtype float16 --ep_size 4 --ep_rank 0 --autotune --use_cuda_events --no_cuda_graph
--routine cutlass_fused_moe --num_tokens 1 --hidden_size 5120 --intermediate_size 8192 --num_experts 128 --top_k 1 --cutlass_variant base --input_dtype float16 --tp_size 4 --tp_rank 0 --autotune --use_cuda_events --no_cuda_graph
--routine cutlass_fused_moe --num_tokens 128 --hidden_size 5120 --intermediate_size 8192 --num_experts 128 --top_k 1 --cutlass_variant base --input_dtype float16 --autotune --use_cuda_events --no_cuda_graph
--routine cutlass_fused_moe --num_tokens 128 --hidden_size 5120 --intermediate_size 8192 --num_experts 128 --top_k 1 --cutlass_variant fp8 --input_dtype float16 --autotune --use_cuda_events --no_cuda_graph
--routine cutlass_fused_moe --num_tokens 128 --hidden_size 5120 --intermediate_size 8192 --num_experts 128 --top_k 1 --cutlass_variant nvfp4 --input_dtype float16 --autotune --use_cuda_events --no_cuda_graph
--routine cutlass_fused_moe --num_tokens 128 --hidden_size 5120 --intermediate_size 8192 --num_experts 128 --top_k 1 --cutlass_variant nvfp4 --quantized_input --input_dtype float16 --autotune --use_cuda_events --no_cuda_graph
--routine cutlass_fused_moe --num_tokens 128 --hidden_size 5120 --intermediate_size 8192 --num_experts 128 --top_k 1 --cutlass_variant base --input_dtype float16 --ep_size 4 --ep_rank 0 --autotune --use_cuda_events --no_cuda_graph
--routine cutlass_fused_moe --num_tokens 128 --hidden_size 5120 --intermediate_size 8192 --num_experts 128 --top_k 1 --cutlass_variant base --input_dtype float16 --tp_size 4 --tp_rank 0 --autotune --use_cuda_events --no_cuda_graph
--routine cutlass_fused_moe --num_tokens 1024 --hidden_size 5120 --intermediate_size 8192 --num_experts 128 --top_k 1 --cutlass_variant base --input_dtype float16 --autotune --use_cuda_events --no_cuda_graph
--routine cutlass_fused_moe --num_tokens 1024 --hidden_size 5120 --intermediate_size 8192 --num_experts 128 --top_k 1 --cutlass_variant fp8 --input_dtype float16 --autotune --use_cuda_events --no_cuda_graph
--routine cutlass_fused_moe --num_tokens 1024 --hidden_size 5120 --intermediate_size 8192 --num_experts 128 --top_k 1 --cutlass_variant nvfp4 --input_dtype float16 --autotune --use_cuda_events --no_cuda_graph
--routine cutlass_fused_moe --num_tokens 1024 --hidden_size 5120 --intermediate_size 8192 --num_experts 128 --top_k 1 --cutlass_variant nvfp4 --quantized_input --input_dtype float16 --autotune --use_cuda_events --no_cuda_graph
--routine cutlass_fused_moe --num_tokens 1024 --hidden_size 5120 --intermediate_size 8192 --num_experts 128 --top_k 1 --cutlass_variant base --input_dtype float16 --ep_size 4 --ep_rank 0 --autotune --use_cuda_events --no_cuda_graph
--routine cutlass_fused_moe --num_tokens 1024 --hidden_size 5120 --intermediate_size 8192 --num_experts 128 --top_k 1 --cutlass_variant base --input_dtype float16 --tp_size 4 --tp_rank 0 --autotune --use_cuda_events --no_cuda_graph
--routine cutlass_fused_moe --num_tokens 1 --hidden_size 7168 --intermediate_size 2048 --num_experts 256 --top_k 8 --cutlass_variant base --input_dtype float16 --autotune --use_cuda_events --no_cuda_graph
--routine cutlass_fused_moe --num_tokens 1 --hidden_size 7168 --intermediate_size 2048 --num_experts 256 --top_k 8 --cutlass_variant fp8 --input_dtype float16 --autotune --use_cuda_events --no_cuda_graph
--routine cutlass_fused_moe --num_tokens 1 --hidden_size 7168 --intermediate_size 2048 --num_experts 256 --top_k 8 --cutlass_variant nvfp4 --input_dtype float16 --autotune --use_cuda_events --no_cuda_graph
--routine cutlass_fused_moe --num_tokens 1 --hidden_size 7168 --intermediate_size 2048 --num_experts 256 --top_k 8 --cutlass_variant nvfp4 --quantized_input --input_dtype float16 --autotune --use_cuda_events --no_cuda_graph
--routine cutlass_fused_moe --num_tokens 1 --hidden_size 7168 --intermediate_size 2048 --num_experts 256 --top_k 8 --cutlass_variant base --input_dtype float16 --ep_size 4 --ep_rank 0 --autotune --use_cuda_events --no_cuda_graph
--routine cutlass_fused_moe --num_tokens 1 --hidden_size 7168 --intermediate_size 2048 --num_experts 256 --top_k 8 --cutlass_variant base --input_dtype float16 --tp_size 4 --tp_rank 0 --autotune --use_cuda_events --no_cuda_graph
--routine cutlass_fused_moe --num_tokens 128 --hidden_size 7168 --intermediate_size 2048 --num_experts 256 --top_k 8 --cutlass_variant base --input_dtype float16 --autotune --use_cuda_events --no_cuda_graph
--routine cutlass_fused_moe --num_tokens 128 --hidden_size 7168 --intermediate_size 2048 --num_experts 256 --top_k 8 --cutlass_variant fp8 --input_dtype float16 --autotune --use_cuda_events --no_cuda_graph
--routine cutlass_fused_moe --num_tokens 128 --hidden_size 7168 --intermediate_size 2048 --num_experts 256 --top_k 8 --cutlass_variant nvfp4 --input_dtype float16 --autotune --use_cuda_events --no_cuda_graph
--routine cutlass_fused_moe --num_tokens 128 --hidden_size 7168 --intermediate_size 2048 --num_experts 256 --top_k 8 --cutlass_variant nvfp4 --quantized_input --input_dtype float16 --autotune --use_cuda_events --no_cuda_graph
--routine cutlass_fused_moe --num_tokens 128 --hidden_size 7168 --intermediate_size 2048 --num_experts 256 --top_k 8 --cutlass_variant base --input_dtype float16 --ep_size 4 --ep_rank 0 --autotune --use_cuda_events --no_cuda_graph
--routine cutlass_fused_moe --num_tokens 128 --hidden_size 7168 --intermediate_size 2048 --num_experts 256 --top_k 8 --cutlass_variant base --input_dtype float16 --tp_size 4 --tp_rank 0 --autotune --use_cuda_events --no_cuda_graph
--routine cutlass_fused_moe --num_tokens 1024 --hidden_size 7168 --intermediate_size 2048 --num_experts 256 --top_k 8 --cutlass_variant base --input_dtype float16 --autotune --use_cuda_events --no_cuda_graph
--routine cutlass_fused_moe --num_tokens 1024 --hidden_size 7168 --intermediate_size 2048 --num_experts 256 --top_k 8 --cutlass_variant fp8 --input_dtype float16 --autotune --use_cuda_events --no_cuda_graph
--routine cutlass_fused_moe --num_tokens 1024 --hidden_size 7168 --intermediate_size 2048 --num_experts 256 --top_k 8 --cutlass_variant nvfp4 --input_dtype float16 --autotune --use_cuda_events --no_cuda_graph
--routine cutlass_fused_moe --num_tokens 1024 --hidden_size 7168 --intermediate_size 2048 --num_experts 256 --top_k 8 --cutlass_variant nvfp4 --quantized_input --input_dtype float16 --autotune --use_cuda_events --no_cuda_graph
--routine cutlass_fused_moe --num_tokens 1024 --hidden_size 7168 --intermediate_size 2048 --num_experts 256 --top_k 8 --cutlass_variant base --input_dtype float16 --ep_size 4 --ep_rank 0 --autotune --use_cuda_events --no_cuda_graph
--routine cutlass_fused_moe --num_tokens 1024 --hidden_size 7168 --intermediate_size 2048 --num_experts 256 --top_k 8 --cutlass_variant base --input_dtype float16 --tp_size 4 --tp_rank 0 --autotune --use_cuda_events --no_cuda_graph
--routine cutlass_fused_moe --num_tokens 1 --hidden_size 2048 --intermediate_size 768 --num_experts 128 --top_k 8 --cutlass_variant nvfp4 --quantized_input --input_dtype bfloat16 --activation-type Swiglu --tp_size 1 --tp_rank 0 --ep_size 1 --ep_rank 0 --use_cuda_events --no_cuda_graph
--routine cutlass_fused_moe --num_tokens 128 --hidden_size 2048 --intermediate_size 768 --num_experts 128 --top_k 8 --cutlass_variant nvfp4 --quantized_input --input_dtype bfloat16 --activation-type Swiglu --tp_size 1 --tp_rank 0 --ep_size 1 --ep_rank 0 --use_cuda_events --no_cuda_graph
--routine cutlass_fused_moe --num_tokens 1024 --hidden_size 2048 --intermediate_size 768 --num_experts 128 --top_k 8 --cutlass_variant nvfp4 --quantized_input --input_dtype bfloat16 --activation-type Swiglu --tp_size 1 --tp_rank 0 --ep_size 1 --ep_rank 0 --use_cuda_events --no_cuda_graph
--routine cutlass_fused_moe --num_tokens 1 --hidden_size 2048 --intermediate_size 1536 --num_experts 128 --top_k 8 --cutlass_variant nvfp4 --quantized_input --input_dtype bfloat16 --activation-type Swiglu --tp_size 1 --tp_rank 0 --ep_size 1 --ep_rank 0 --use_cuda_events --no_cuda_graph
--routine cutlass_fused_moe --num_tokens 128 --hidden_size 2048 --intermediate_size 1536 --num_experts 128 --top_k 8 --cutlass_variant nvfp4 --quantized_input --input_dtype bfloat16 --activation-type Swiglu --tp_size 1 --tp_rank 0 --ep_size 1 --ep_rank 0 --use_cuda_events --no_cuda_graph
--routine cutlass_fused_moe --num_tokens 1024 --hidden_size 2048 --intermediate_size 1536 --num_experts 128 --top_k 8 --cutlass_variant nvfp4 --quantized_input --input_dtype bfloat16 --activation-type Swiglu --tp_size 1 --tp_rank 0 --ep_size 1 --ep_rank 0 --use_cuda_events --no_cuda_graph

Click to view perf data on RTX Pro 6000 SM120) and Spark (SM121)

Perf data demonstrated that untargeted cases have no performance impact. Targeted (NVFP4) CUTLASS MoE cases see speedup on decode cases.

Summary table:

Category	Count	Spark Geomean	RTX P6K Geomean
Targeted NVFP4 (decode-optimized)	6	1.051x	1.021x
Unaffected (base/fp8/large batch)	72	1.008x	1.006x
All cases	78	1.011x	1.007x

Full table:

#	Test Case	Spark Before (ms)	Spark After (ms)	Spark Speedup	RTX P6K Before (ms)	RTX P6K After (ms)	RTX P6K Speedup	Category
1	T=1 H=5120 I=8192 E=128 k=1 base	1.3061	1.3256	0.985x	0.3041	0.3033	1.003x
2	T=1 H=5120 I=8192 E=128 k=1 fp8	0.8187	0.8264	0.991x	0.2152	0.2153	0.999x
3	T=1 H=5120 I=8192 E=128 k=1 nvfp4	0.4329	0.4157	1.041x	0.1300	0.1331	0.977x
4	T=1 H=5120 I=8192 E=128 k=1 nvfp4qi	0.4335	0.4127	1.051x	0.1311	0.1320	0.993x
5	T=1 H=5120 I=8192 E=128 k=1 base ep4	0.0144	0.0152	0.950x	0.0154	0.0142	1.086x
6	T=1 H=5120 I=8192 E=128 k=1 base tp4	0.3676	0.3757	0.978x	0.1817	0.1833	0.991x
7	T=128 H=5120 I=8192 E=128 k=1 base	93.9090	91.7755	1.023x	12.7980	12.8000	1.000x
8	T=128 H=5120 I=8192 E=128 k=1 fp8	55.0527	54.1031	1.018x	6.5404	6.5444	0.999x
9	T=128 H=5120 I=8192 E=128 k=1 nvfp4	25.2912	25.2012	1.004x	3.7753	3.7632	1.003x
10	T=128 H=5120 I=8192 E=128 k=1 nvfp4qi	25.1966	25.2118	0.999x	3.7776	3.7699	1.002x
11	T=128 H=5120 I=8192 E=128 k=1 base ep4	22.4435	22.5981	0.993x	3.1017	3.1058	0.999x
12	T=128 H=5120 I=8192 E=128 k=1 base tp4	22.4918	22.6853	0.991x	3.2777	3.2813	0.999x
13	T=1024 H=5120 I=8192 E=128 k=1 base	144.5477	145.0290	0.997x	21.3314	21.3340	1.000x
14	T=1024 H=5120 I=8192 E=128 k=1 fp8	82.2042	81.8906	1.004x	11.0048	11.0082	1.000x
15	T=1024 H=5120 I=8192 E=128 k=1 nvfp4	39.7629	39.8295	0.998x	6.3147	6.2995	1.002x
16	T=1024 H=5120 I=8192 E=128 k=1 nvfp4qi	39.6928	39.8489	0.996x	6.3218	6.2758	1.007x
17	T=1024 H=5120 I=8192 E=128 k=1 base ep4	36.2588	36.4861	0.994x	5.4456	5.4447	1.000x
18	T=1024 H=5120 I=8192 E=128 k=1 base tp4	35.4185	35.4508	0.999x	5.4314	5.4313	1.000x
19	T=1 H=7168 I=2048 E=256 k=8 base	3.2219	3.2369	0.995x	0.8597	0.8597	1.000x
20	T=1 H=7168 I=2048 E=256 k=8 fp8	1.8355	1.8222	1.007x	0.3420	0.3417	1.001x
21	T=1 H=7168 I=2048 E=256 k=8 nvfp4	1.0445	0.9912	1.054x	0.2591	0.2437	1.063x
22	T=1 H=7168 I=2048 E=256 k=8 nvfp4qi	1.0440	0.9820	1.063x	0.2591	0.2448	1.059x
23	T=1 H=7168 I=2048 E=256 k=8 base ep4	1.2269	1.2545	0.978x	0.3676	0.3670	1.001x
24	T=1 H=7168 I=2048 E=256 k=8 base tp4	0.8725	0.8841	0.987x	0.2396	0.2396	1.000x
25	T=128 H=7168 I=2048 E=256 k=8 base	97.3592	97.6644	0.997x	14.6278	14.6295	1.000x
26	T=128 H=7168 I=2048 E=256 k=8 fp8	52.8793	52.3160	1.011x	7.4488	7.4495	1.000x
27	T=128 H=7168 I=2048 E=256 k=8 nvfp4	27.4207	27.1848	1.009x	4.3588	4.3220	1.009x
28	T=128 H=7168 I=2048 E=256 k=8 nvfp4qi	27.5014	27.1511	1.013x	4.3604	4.3237	1.008x
29	T=128 H=7168 I=2048 E=256 k=8 base ep4	23.9165	23.8423	1.003x	3.7266	3.7261	1.000x
30	T=128 H=7168 I=2048 E=256 k=8 base tp4	24.0706	24.1821	0.995x	3.7711	3.7720	1.000x
31	T=1024 H=7168 I=2048 E=256 k=8 base	126.3262	127.5760	0.990x	17.5073	17.4586	1.003x
32	T=1024 H=7168 I=2048 E=256 k=8 fp8	68.2070	68.9878	0.989x	8.3225	8.3213	1.000x
33	T=1024 H=7168 I=2048 E=256 k=8 nvfp4	30.1287	30.1583	0.999x	4.8560	4.8396	1.003x
34	T=1024 H=7168 I=2048 E=256 k=8 nvfp4qi	29.9361	29.7770	1.005x	4.8430	4.8278	1.003x
35	T=1024 H=7168 I=2048 E=256 k=8 base ep4	30.0261	29.8761	1.005x	4.1991	4.1916	1.002x
36	T=1024 H=7168 I=2048 E=256 k=8 base tp4	29.5162	29.1558	1.012x	4.2802	4.2811	1.000x
37	T=1 H=5120 I=8192 E=128 k=1 base	1.2212	1.2314	0.992x	0.2519	0.2519	1.000x
38	T=1 H=5120 I=8192 E=128 k=1 fp8	0.6404	0.6404	1.000x	0.1475	0.1484	0.994x
39	T=1 H=5120 I=8192 E=128 k=1 nvfp4	0.4505	0.4197	1.073x	0.1252	0.1230	1.018x
40	T=1 H=5120 I=8192 E=128 k=1 nvfp4qi	0.4541	0.4108	1.106x	0.1251	0.1215	1.030x
41	T=1 H=5120 I=8192 E=128 k=1 base ep4	0.0175	0.0174	1.006x	0.0132	0.0123	1.070x
42	T=1 H=5120 I=8192 E=128 k=1 base tp4	0.3666	0.3687	0.994x	0.0901	0.0891	1.011x
43	T=128 H=5120 I=8192 E=128 k=1 base	90.6373	91.2210	0.994x	12.7805	12.7817	1.000x
44	T=128 H=5120 I=8192 E=128 k=1 fp8	44.6182	44.7585	0.997x	6.4339	6.4344	1.000x
45	T=128 H=5120 I=8192 E=128 k=1 nvfp4	25.6261	25.4186	1.008x	3.7643	3.7560	1.002x
46	T=128 H=5120 I=8192 E=128 k=1 nvfp4qi	25.6532	25.3584	1.012x	3.7627	3.7540	1.002x
47	T=128 H=5120 I=8192 E=128 k=1 base ep4	22.2766	21.8055	1.022x	3.1002	3.0995	1.000x
48	T=128 H=5120 I=8192 E=128 k=1 base tp4	22.5305	22.3909	1.006x	3.2462	3.2461	1.000x
49	T=1024 H=5120 I=8192 E=128 k=1 base	140.0064	139.6009	1.003x	21.2091	21.2116	1.000x
50	T=1024 H=5120 I=8192 E=128 k=1 fp8	69.0247	69.4502	0.994x	10.7856	10.8737	0.992x
51	T=1024 H=5120 I=8192 E=128 k=1 nvfp4	40.2212	39.9195	1.008x	6.2956	6.2756	1.003x
52	T=1024 H=5120 I=8192 E=128 k=1 nvfp4qi	40.2185	39.9197	1.007x	6.3109	6.2855	1.004x
53	T=1024 H=5120 I=8192 E=128 k=1 base ep4	35.2440	34.7130	1.015x	5.3981	5.3986	1.000x
54	T=1024 H=5120 I=8192 E=128 k=1 base tp4	35.0929	34.7923	1.009x	5.3969	5.3954	1.000x
55	T=1 H=7168 I=2048 E=256 k=8 base	3.1447	3.1472	0.999x	0.5619	0.5993	0.938x
56	T=1 H=7168 I=2048 E=256 k=8 fp8	1.6363	1.6096	1.017x	0.3256	0.3256	1.000x
57	T=1 H=7168 I=2048 E=256 k=8 nvfp4	1.0844	1.0019	1.082x	0.2468	0.2355	1.048x
58	T=1 H=7168 I=2048 E=256 k=8 nvfp4qi	1.0869	1.0040	1.083x	0.2499	0.2336	1.069x
59	T=1 H=7168 I=2048 E=256 k=8 base ep4	1.2902	1.2902	1.000x	0.1956	0.1956	1.000x
60	T=1 H=7168 I=2048 E=256 k=8 base tp4	0.8631	0.8637	0.999x	0.2005	0.1997	1.004x
61	T=128 H=7168 I=2048 E=256 k=8 base	95.4863	95.3204	1.002x	14.5634	14.5612	1.000x
62	T=128 H=7168 I=2048 E=256 k=8 fp8	46.0011	45.9304	1.002x	7.3332	7.3367	1.000x
63	T=128 H=7168 I=2048 E=256 k=8 nvfp4	27.3023	26.9143	1.014x	4.3172	4.2922	1.006x
64	T=128 H=7168 I=2048 E=256 k=8 nvfp4qi	27.4191	26.9065	1.019x	4.3254	4.2943	1.007x
65	T=128 H=7168 I=2048 E=256 k=8 base ep4	24.3881	24.2863	1.004x	3.7139	3.7150	1.000x
66	T=128 H=7168 I=2048 E=256 k=8 base tp4	23.5841	23.6456	0.997x	3.7500	3.7499	1.000x
67	T=1024 H=7168 I=2048 E=256 k=8 base	104.3629	104.1310	1.002x	15.7031	15.7159	0.999x
68	T=1024 H=7168 I=2048 E=256 k=8 fp8	49.5421	49.4592	1.002x	7.9900	7.9933	1.000x
69	T=1024 H=7168 I=2048 E=256 k=8 nvfp4	30.2330	29.8347	1.013x	4.7944	4.7812	1.003x
70	T=1024 H=7168 I=2048 E=256 k=8 nvfp4qi	29.9776	29.6689	1.010x	4.7790	4.7718	1.002x
71	T=1024 H=7168 I=2048 E=256 k=8 base ep4	26.2313	26.3157	0.997x	3.9859	3.9855	1.000x
72	T=1024 H=7168 I=2048 E=256 k=8 base tp4	27.7164	27.7785	0.998x	4.2291	4.2261	1.001x
73	T=1 H=2048 I=768 E=128 k=8 nvfp4qi	0.2232	0.1987	1.123x	0.0705	0.0686	1.028x	TARGETED
74	T=128 H=2048 I=768 E=128 k=8 nvfp4qi	1.6925	1.6194	1.045x	0.3602	0.3553	1.014x	TARGETED
75	T=1024 H=2048 I=768 E=128 k=8 nvfp4qi	2.2461	2.2195	1.012x	0.4055	0.3941	1.029x	TARGETED
76	T=1 H=2048 I=1536 E=128 k=8 nvfp4qi	0.3179	0.2878	1.105x	0.0967	0.0886	1.092x	TARGETED
77	T=128 H=2048 I=1536 E=128 k=8 nvfp4qi	3.1408	3.1017	1.013x	0.5934	0.5847	1.015x	TARGETED
78	T=1024 H=2048 I=1536 E=128 k=8 nvfp4qi	4.0100	3.9479	1.016x	0.6842	0.6788	1.008x	TARGETED

🔍 Related Issues

#3013

🚀 Pull Request Checklist

Thank you for contributing to FlashInfer! Before we review your pull request, please make sure the following items are complete.

✅ Pre-commit Checks

I have installed pre-commit by running pip install pre-commit (or used your preferred method).
I have installed the hooks with pre-commit install.
I have run the hooks manually with pre-commit run --all-files and fixed any reported issues.

If you are unsure about how to set up pre-commit, see the pre-commit documentation.

🧪 Tests

Tests have been added or updated as needed.
All tests are passing (unittest, etc.).

Reviewer Notes

Summary by CodeRabbit

Performance
- Reduced unnecessary work for empty experts and tightened activation/grid sizing so only required entries are processed, improving memory efficiency and GPU utilization.
Bug Fixes
- Fixed edge cases that could produce incorrect outputs when many experts are empty and removed unsafe padding work that risked uninitialized values.
Tests
- Added CUDA-architecture gated tests verifying numerical correctness and memory-safety (including poisoning checks) for the modified paths.

coderabbitai · 2026-04-08T04:50:00Z

Note

Reviews paused

It looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the reviews.auto_review.auto_pause_after_reviewed_commits setting.

Use the following commands to manage reviews:

@coderabbitai resume to resume automatic reviews.
@coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

▶️ Resume reviews
🔍 Trigger review

📝 Walkthrough

Walkthrough

Replaced inline PTX grid-dependency intrinsics with CUDA runtime helpers, added early-exit for experts with zero GEMM work (still triggering programmatic-launch completion), removed N-dimension SF padding loops, adjusted kernel launch sizing and thread counts, and added CUDA-gated tests covering the removed N-dim padding behavior.

Changes

Cohort / File(s)	Summary
Cutlass fused-MoE kernels `csrc/fused_moe/cutlass_backend/cutlass_fused_moe_kernels.cuh`	Replaced inline PTX grid-dependency calls with `cudaGridDependencySynchronize()` / `cudaTriggerProgrammaticLaunchCompletion()`. Added early-return when `gemm_m == 0` (still triggers programmatic-launch completion on CUDA_ARCH>=900). Removed N-dimension SF padding loops from `expandInputRowsKernel` and `doActivationKernel`; preserved K-dimension zeroing along valid token path. Removed padding-driven block accounting from launchers and added `__launch_bounds__` to activation kernel. Reduced per-expert stride-setup threads from `min(1024, ...)` to `min(32, ...)` and recomputed blocks.
MoE tests `tests/moe/test_trtllm_cutlass_fused_moe.py`	Added CUDA-only parametrized tests `test_moe_nvfp4_ndim_padding_safety` and `test_moe_mxfp8_mxfp4_ndim_padding_safety` (with module-level param constants) to verify correctness and to detect any uninitialized SF usage after removing N-dim padding rows. Includes quantized/unquantized paths and GPU-memory poisoning checks.

Sequence Diagram(s)

sequenceDiagram
    participant Host as Host
    participant Runner as CutlassMoeFCRunner
    participant Kernel as CUDA Kernel
    participant GridDep as GridDep Runtime

    Host->>Runner: prepare params (expanded_num_tokens, num_experts_per_node, k, ...)
    Runner->>Kernel: launch kernel (expandInputRows / doActivation / computeStrides)
    Kernel->>Kernel: evaluate per-expert tokens_to_expert
    alt tokens_to_expert == 0
        Kernel->>GridDep: call runtime trigger (cudaTriggerProgrammaticLaunchCompletion)
        Kernel-->>Runner: early return (skip per-expert setup)
    else tokens present
        Kernel->>Kernel: perform per-token K-dim zeroing, stride/pointer setup, compute strides
        Kernel-->>Runner: complete writes
    end
    Runner-->>Host: report completion

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~45 minutes

Possibly related PRs

[NVIDIA] fix(jit): enable GDC for CUTLASS fused MoE PDL — prevent random crashes on SM12x #2913: touches fused MoE CUTLASS grid-dependency control and inline GDC usage that relates to the grid-dependency runtime change here.

Suggested reviewers

aleozlx
samuellees
IwakuraRein
jiahanc
nv-yunzheq

Poem

🐰
I hop through kernels, light and keen,
Trimming padding from each hidden scene,
I nudge the grids to wait then go,
Threads shrink where idle winds would blow,
A tiny nibble, tidy and clean.

🚥 Pre-merge checks | ✅ 3

✅ Passed checks (3 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly and specifically summarizes the main objective of the PR: optimizing CUTLASS MoE helper kernels for small-batch decode workloads.
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.
Description check	✅ Passed	The PR description is comprehensive and well-structured with clear sections covering motivation, changes, correctness validation, and performance data.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create PR with unit tests

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

gemini-code-assist

Code Review

This pull request optimizes Fused MoE kernels by skipping expensive setup for experts with no assigned tokens and refactoring padding operations to be expert-driven and flattened for better thread utilization. Review feedback highlighted a critical unit mismatch in the merged activation loop for FP4/MXFP8, where indexing units for loop bounds and element access are inconsistent, leading to incomplete processing. Additionally, an early exit in the stride computation kernel was flagged for potentially skipping required initialization of problem shapes for experts with zero tokens, which could cause undefined behavior in CUTLASS.

gemini-code-assist · 2026-04-08T04:51:30Z

+    int64_t const loop_elems = (IsNVFP4 || IsMXFP8)
+        ? padded_inter_size / VecSize  // cover both real elements and K-dim padding
+        : num_elems_in_col;
+    bool const do_k_padding = (IsNVFP4 || IsMXFP8) && (padded_inter_size > inter_size);
+
    ActFn fn{};
    fn.alpha = gate_alpha;
    fn.beta = gate_beta;
    fn.limit = gate_limit;
-    for (int64_t elem_index = start_offset; elem_index < num_elems_in_col; elem_index += stride) {
+    for (int64_t elem_index = start_offset; elem_index < loop_elems; elem_index += stride) {
+      // K-dim padding region: write zero SF only (no activation compute)
+      if (do_k_padding && elem_index >= num_elems_in_col) {
+        writeSF<VecSize, VecSize>(num_tokens_before_expert, expert, /*source_row*/ -1, token,
+                                  elem_index, padded_inter_size, fc2_act_sf_flat,
+                                  /* input_sf */ nullptr);
+        continue;
+      }
+
      auto fc1_value =
          arrayConvert<GemmResultElem, ComputeElem>(gemm_result_vec[elem_index + gated_off_vec]);
      if (bias_ptr) {


There is a critical unit mismatch in the merged activation loop for FP4/MXFP8. loop_elems is calculated in units of VecSize (16 or 32), but the loop body indexes gemm_result_vec and output_vec using elem_index as if it were in units of ACTIVATION_ELEM_PER_THREAD (8).

For example, if inter_size is 1024 and VecSize is 16, loop_elems will be 64. The loop will only iterate 64 times, processing elements 0-511 (since gemm_result_vec[elem_index] accesses 8 elements), leaving the second half of the columns unprocessed. Furthermore, the padding check elem_index >= num_elems_in_col compares values with different units (units of VecSize vs units of 8), leading to incorrect branching between compute and padding logic. This logic needs to be corrected to ensure all elements are processed and indexing is consistent.

Good catch — you're absolutely right. The main loop indexes in units of ACTIVATION_ELEM_PER_THREAD (8) while the K-dim padding loop indexes in units of VecSize (16 for NVFP4). Merging them into a single loop with padded_inter_size / VecSize as the upper bound causes the main loop to under-iterate (48 iterations instead of 96 for inter_size=768), leaving half the columns unprocessed. The elem_index >= num_elems_in_col check also compares values in different units.

Reverted to the original separate K-dim padding loop in the next commit. In practice, this merge was a no-op for the target workload anyway — Qwen3-30B-A3B has inter_size=768 which is already 128-aligned, so the K-dim padding loop executes zero iterations. The actual performance wins in this PR come from the N-dim padding restructuring and computeStrides changes, which are unaffected.

gemini-code-assist · 2026-04-08T04:51:30Z

+  if (gemm_m == 0) {
+#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
+    asm volatile("griddepcontrol.launch_dependents;");
+#endif
+    return;
+  }


The early exit for gemm_m == 0 is placed before the initialization of int4_groupwise_params.shape.problem_shapes (lines 1273-1285). CUTLASS grouped GEMM typically requires all problem shapes to be initialized, even for zero-M problems, to correctly traverse the problem list. Skipping this initialization for experts with no tokens may lead to undefined behavior or illegal memory accesses if the visitor reads uninitialized data from the problem shape buffer. Consider moving the early exit after all problem shape assignments.

Good point — moved the early exit after all problem shape assignments (both regular and int4_groupwise). CUTLASS's tile scheduler may read problem shapes for all experts to traverse the problem list, so they need to be fully initialized even for zero-M entries. The remaining work after the early exit (alpha scales, block scaling factors, strides, pointers) is only consumed by CUTLASS for experts it actually processes, so skipping those for zero-M experts is safe.

coderabbitai

Actionable comments posted: 1

🤖 Prompt for all review comments with AI agents

Verify each finding against the current code and only fix it if needed.

Inline comments:
In `@csrc/fused_moe/cutlass_backend/cutlass_fused_moe_kernels.cuh`:
- Around line 2165-2182: The loop currently computes loop_elems in SF-block
units (loop_elems = padded_inter_size / VecSize) while the body still indexes
activation chunks using num_elems_in_col (which is in ACTIVATION_ELEM_PER_THREAD
units), so most activation chunks are skipped and the do_k_padding branch is
never hit; fix this by making the loop iterate in ACTIVATION_ELEM_PER_THREAD
units (e.g., set loop_elems = padded_inter_size / ACTIVATION_ELEM_PER_THREAD or
restore the separate K-padding loop) and only translate indices to SF-block
units when calling writeSF; update use of loop_elems, num_elems_in_col,
do_k_padding, and the writeSF call so that activation compute and K-padding
checks use the same unit (ACTIVATION_ELEM_PER_THREAD) and convert elem_index to
VecSize/SF-block index only for writeSF/padding writes.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

Push a commit to this branch (recommended)
Create a new PR with the fixes

ℹ️ Review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: daee07d9-476c-4eb9-b740-4a803176e586

📥 Commits

Reviewing files that changed from the base of the PR and between 1a2c2eb and e9224bf.

📒 Files selected for processing (1)

csrc/fused_moe/cutlass_backend/cutlass_fused_moe_kernels.cuh

coderabbitai

Actionable comments posted: 3

🤖 Prompt for all review comments with AI agents

Verify each finding against the current code and only fix it if needed.

Inline comments:
In `@tests/moe/test_trtllm_cutlass_fused_moe.py`:
- Around line 2122-2125: The test uses torch.cuda.get_device_capability()
directly to gate NVFP4 skips; replace that direct query with the repo skip
helpers from flashinfer.utils to keep gating consistent: use
get_compute_capability() or the specific predicates like
is_sm100a_supported(...) / is_sm12x_supported(...) (and is_sm90a_supported if
relevant) instead of torch.cuda.get_device_capability(), updating the two skip
decorators (around the blocks shown and the similar one at ~2276-2279) to call
the appropriate helper so the tests follow the established pattern.
- Around line 2188-2213: The test should deterministically poison the temp/SF
scratch path so stale reads fail reliably: before calling
fused_moe.cutlass_fused_moe, fill the output buffer (flash_output) with a
sentinel (e.g., NaN or 0xFF) and, when quantized_input is true, explicitly
allocate and pass an input SF/scratch tensor (the input_sf parameter) filled
with the same sentinel instead of leaving it None; ensure the quant_scales list
and fp4_quantize usage remain unchanged but drive the kernel to use the provided
input_sf so the kernel’s internal SF scratch path is exercised and any
stale-data bug is exposed.
- Line 2148: The inline lambda assigned to round_up causes an E731 violation;
remove that lambda and replace any uses of round_up with the existing ceil_div()
helper already defined in this test file (i.e., delete the line defining
round_up and call ceil_div(x, y) wherever round_up(x, y) was used so behavior
remains identical).

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

Push a commit to this branch (recommended)
Create a new PR with the fixes

ℹ️ Review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: 4916041b-893d-4dbc-8d5c-dd270773c319

📥 Commits

Reviewing files that changed from the base of the PR and between a33b48c and 37dc223.

📒 Files selected for processing (2)

csrc/fused_moe/cutlass_backend/cutlass_fused_moe_kernels.cuh
tests/moe/test_trtllm_cutlass_fused_moe.py

coderabbitai · 2026-04-09T04:57:36Z

+@pytest.mark.skipif(
+    torch.cuda.get_device_capability()[0] not in [10, 11, 12],
+    reason="NVFP4 is only supported on SM100, SM110 and SM120/SM121",
+)


⚠️ Potential issue | 🟡 Minor

Use the flashinfer.utils architecture helpers for these skip guards.

These new tests query torch.cuda.get_device_capability() directly even though this file already uses the repo’s skip helpers for the same purpose. Please mirror the nearby is_sm100a_supported(...) / is_sm12x_supported(...) pattern so the gating logic stays consistent across the test suite.

♻️ Suggested change

`@pytest.mark.skipif`( - torch.cuda.get_device_capability()[0] not in [10, 11, 12], - reason="NVFP4 is only supported on SM100, SM110 and SM120/SM121", + not is_sm100a_supported(torch.device("cuda")) + and not is_sm12x_supported(torch.device("cuda")), + reason="NVFP4 is only supported on SM100+", )

`@pytest.mark.skipif`( - torch.cuda.get_device_capability()[0] not in [10, 11, 12], - reason="MXFP8xMXFP4 is only supported on SM100, SM110 and SM120/SM121", + not is_sm100a_supported(torch.device("cuda")) + and not is_sm12x_supported(torch.device("cuda")), + reason="MXFP8xMXFP4 is only supported on SM100+", )

As per coding guidelines: Use flashinfer.utils functions (get_compute_capability(), is_sm90a_supported(), is_sm100a_supported()) to skip tests on unsupported GPU architectures.

Also applies to: 2276-2279

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed. In `@tests/moe/test_trtllm_cutlass_fused_moe.py` around lines 2122 - 2125, The test uses torch.cuda.get_device_capability() directly to gate NVFP4 skips; replace that direct query with the repo skip helpers from flashinfer.utils to keep gating consistent: use get_compute_capability() or the specific predicates like is_sm100a_supported(...) / is_sm12x_supported(...) (and is_sm90a_supported if relevant) instead of torch.cuda.get_device_capability(), updating the two skip decorators (around the blocks shown and the similar one at ~2276-2279) to call the appropriate helper so the tests follow the established pattern.

coderabbitai · 2026-04-09T04:57:37Z

+    torch.manual_seed(42)
+    otype = torch.bfloat16
+    quant_blocksize = 16
+    round_up = lambda x, y: (x + y - 1) // y * y


⚠️ Potential issue | 🟡 Minor

🧩 Analysis chain

🏁 Script executed:

# First, check if the file exists find . -name "test_trtllm_cutlass_fused_moe.py" -type f

Repository: flashinfer-ai/flashinfer

Length of output: 112

🏁 Script executed:

# Read the file to check line 2148 and surrounding context wc -l tests/moe/test_trtllm_cutlass_fused_moe.py

Repository: flashinfer-ai/flashinfer

Length of output: 115

🏁 Script executed:

# Get the content around line 2148 sed -n '2140,2155p' tests/moe/test_trtllm_cutlass_fused_moe.py

Repository: flashinfer-ai/flashinfer

Length of output: 523

🏁 Script executed:

# Search for ceil_div import or definition in the file rg "ceil_div" tests/moe/test_trtllm_cutlass_fused_moe.py

Repository: flashinfer-ai/flashinfer

Length of output: 434

🏁 Script executed:

# Check imports at the top of the file head -50 tests/moe/test_trtllm_cutlass_fused_moe.py

Repository: flashinfer-ai/flashinfer

Length of output: 1610

🏁 Script executed:

# Search for ceil_div definition or import in the file rg "from.*ceil_div|import.*ceil_div|def ceil_div" tests/moe/test_trtllm_cutlass_fused_moe.py

Repository: flashinfer-ai/flashinfer

Length of output: 104

🏁 Script executed:

# Check where ceil_div is defined - search in the entire codebase rg "def ceil_div" --type python

Repository: flashinfer-ai/flashinfer

Length of output: 98

🏁 Script executed:

# Let's check the full import section and look for ceil_div more broadly rg "ceil_div" -B 5 -A 5 tests/moe/test_trtllm_cutlass_fused_moe.py | head -50

Repository: flashinfer-ai/flashinfer

Length of output: 1631

Replace the inline lambda with a reusable ceil_div() call to avoid E731 style violation.

Line 2148 assigns a lambda expression, which triggers Ruff E731. The file already defines ceil_div() for ceiling division, so reuse it:

♻️ Suggested change

- round_up = lambda x, y: (x + y - 1) // y * y + def round_up(x: int, y: int) -> int: + return ceil_div(x, y) * y

📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change

round_up = lambda x, y: (x + y - 1) // y * y

def round_up(x: int, y: int) -> int:

return ceil_div(x, y) * y

🧰 Tools

🪛 Ruff (0.15.9)

[error] 2148-2148: Do not assign a lambda expression, use a def

Rewrite round_up as a def

(E731)

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed. In `@tests/moe/test_trtllm_cutlass_fused_moe.py` at line 2148, The inline lambda assigned to round_up causes an E731 violation; remove that lambda and replace any uses of round_up with the existing ceil_div() helper already defined in this test file (i.e., delete the line defining round_up and call ceil_div(x, y) wherever round_up(x, y) was used so behavior remains identical).

coderabbitai · 2026-04-09T04:57:37Z

+    flash_output = torch.zeros_like(x)
+
+    quant_scales = [
+        a1_gs,
+        w1_blockscale.view(torch.int32),
+        1.0 / (a1_gs * w1_gs),
+        a2_gs,
+        w2_blockscale.view(torch.int32),
+        1.0 / (a2_gs * w2_gs),
+    ]
+    hidden_states = x
+    input_sf = None
+    if quantized_input:
+        hidden_states, input_sf = fp4_quantize(x, a1_gs)
+
+    _ = fused_moe.cutlass_fused_moe(
+        hidden_states,
+        selected_experts.to(torch.int),
+        routing_weights,
+        w1_q.contiguous().view(torch.long),
+        w2_q.contiguous().view(torch.long),
+        otype,
+        quant_scales=quant_scales,
+        input_sf=input_sf,
+        output=flash_output,
+    )


⚠️ Potential issue | 🟠 Major

Make the stale-SF regression deterministic.

test_moe_nvfp4_ndim_padding_safety never forces the removed SF-padding rows to contain nonzero garbage, and the MXFP8 variant only proves reuse of a separate 64 MiB block before freeing it again. That still does not guarantee the kernel’s internal SF scratch buffers are poisoned, so a stale-read bug can pass depending on allocator state. Please drive the exact temp-buffer path with a known NaN/0xFF fill or a dedicated debug hook instead of relying on opportunistic caching-allocator reuse.

Also applies to: 2327-2337

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed. In `@tests/moe/test_trtllm_cutlass_fused_moe.py` around lines 2188 - 2213, The test should deterministically poison the temp/SF scratch path so stale reads fail reliably: before calling fused_moe.cutlass_fused_moe, fill the output buffer (flash_output) with a sentinel (e.g., NaN or 0xFF) and, when quantized_input is true, explicitly allocate and pass an input SF/scratch tensor (the input_sf parameter) filled with the same sentinel instead of leaving it None; ensure the quant_scales list and fp4_quantize usage remain unchanged but drive the kernel to use the provided input_sf so the kernel’s internal SF scratch path is exercised and any stale-data bug is exposed.

bkryu · 2026-04-09T06:27:46Z

/bot run

flashinfer-bot · 2026-04-09T06:27:59Z

GitLab MR !524 has been updated with latest changes, and the CI pipeline #48089726 is currently running. I'll report back once the pipeline job completes.

yzh119 · 2026-04-09T18:46:58Z

+  // For decode (1 token, top_k=8, 128 experts), this skips ~120 of 128 experts.
+  if (gemm_m == 0) {
+#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
+    asm volatile("griddepcontrol.launch_dependents;");


Per suggestion in #2558 (comment), we might use cuda native primitives such as cudaTriggerProgrammaticLaunchCompletion here.

Thanks @yzh119, didn't know about the CUDA native APIs. Made substitutions for every PDL PTX in the file in e7bc590 👍

bkryu · 2026-04-09T19:28:54Z

/bot run

flashinfer-bot · 2026-04-09T19:29:15Z

GitLab MR !524 has been updated with latest changes, and the CI pipeline #48144502 is currently running. I'll report back once the pipeline job completes.

djns99 · 2026-04-10T00:25:11Z

+
+    flash_output = torch.zeros_like(x)
+
+    # Poison GPU memory with 0xFF — same approach as the NVFP4 test.


I don't see this logic in nvfp4 test?
And can you confirm that if you remove the K dim padding this test catches it?

I don't see this logic in nvfp4 test?

Sorry that was a stale poisoning attempt. Removed in the latest commit.

I don't see this logic in nvfp4 test?
And can you confirm that if you remove the K dim padding this test catches it?

Yes I can confirm that with the after disapling the K-dim padding and poisoning with 0xFF in commit 6f36003, I see

$ pytest tests/moe/test_trtllm_cutlass_fused_moe.py::test_moe_nvfp4_unaligned_hidden_size -v --tb=short ============================================================================================ test session starts ============================================================================================ ... tests/moe/test_trtllm_cutlass_fused_moe.py::test_moe_nvfp4_unaligned_hidden_size[swiglu-False-otype0-wtype0-2-2-288-128-1] FAILED [ 10%] tests/moe/test_trtllm_cutlass_fused_moe.py::test_moe_nvfp4_unaligned_hidden_size[swiglu-False-otype0-wtype0-2-2-288-128-4] FAILED [ 20%] tests/moe/test_trtllm_cutlass_fused_moe.py::test_moe_nvfp4_unaligned_hidden_size[swiglu-False-otype0-wtype0-2-2-288-192-1] FAILED [ 30%] tests/moe/test_trtllm_cutlass_fused_moe.py::test_moe_nvfp4_unaligned_hidden_size[swiglu-False-otype0-wtype0-2-2-288-192-4] FAILED [ 40%] tests/moe/test_trtllm_cutlass_fused_moe.py::test_moe_nvfp4_unaligned_hidden_size[swiglu-False-otype0-wtype0-2-2-160-192-1] FAILED [ 50%] tests/moe/test_trtllm_cutlass_fused_moe.py::test_moe_nvfp4_unaligned_hidden_size[swiglu-False-otype0-wtype0-2-2-160-192-4] FAILED [ 60%] tests/moe/test_trtllm_cutlass_fused_moe.py::test_moe_nvfp4_unaligned_hidden_size[swiglu-False-otype0-wtype0-2-2-320-160-1] FAILED [ 70%] tests/moe/test_trtllm_cutlass_fused_moe.py::test_moe_nvfp4_unaligned_hidden_size[swiglu-False-otype0-wtype0-2-2-320-160-4] FAILED [ 80%] tests/moe/test_trtllm_cutlass_fused_moe.py::test_moe_nvfp4_unaligned_hidden_size[swiglu-False-otype0-wtype0-2-2-256-192-1] PASSED [ 90%] tests/moe/test_trtllm_cutlass_fused_moe.py::test_moe_nvfp4_unaligned_hidden_size[swiglu-False-otype0-wtype0-2-2-256-192-4] PASSED [100%] ================================================================================================= FAILURES ================================================================================================== ______________________________________________________________ test_moe_nvfp4_unaligned_hidden_size[swiglu-False-otype0-wtype0-2-2-288-128-1] _______________________________________________________________ tests/moe/test_trtllm_cutlass_fused_moe.py:2093: in test_moe_nvfp4_unaligned_hidden_size torch.testing.assert_close(ref_output, flash_output, rtol=2e-1, atol=2e-1) E AssertionError: Tensor-likes are not close! E E Mismatched elements: 288 / 288 (100.0%) E Greatest absolute difference: nan at index (0, 0) (up to 0.2 allowed) E Greatest relative difference: nan at index (0, 0) (up to 0.2 allowed) ______________________________________________________________ test_moe_nvfp4_unaligned_hidden_size[swiglu-False-otype0-wtype0-2-2-288-128-4] _______________________________________________________________ tests/moe/test_trtllm_cutlass_fused_moe.py:2093: in test_moe_nvfp4_unaligned_hidden_size torch.testing.assert_close(ref_output, flash_output, rtol=2e-1, atol=2e-1) E AssertionError: Tensor-likes are not close! E E Mismatched elements: 1152 / 1152 (100.0%) E Greatest absolute difference: nan at index (0, 0) (up to 0.2 allowed) E Greatest relative difference: nan at index (0, 0) (up to 0.2 allowed) ... ... ========================================================================================== short test summary info ========================================================================================== FAILED tests/moe/test_trtllm_cutlass_fused_moe.py::test_moe_nvfp4_unaligned_hidden_size[swiglu-False-otype0-wtype0-2-2-288-128-1] - AssertionError: Tensor-likes are not close! FAILED tests/moe/test_trtllm_cutlass_fused_moe.py::test_moe_nvfp4_unaligned_hidden_size[swiglu-False-otype0-wtype0-2-2-288-128-4] - AssertionError: Tensor-likes are not close! FAILED tests/moe/test_trtllm_cutlass_fused_moe.py::test_moe_nvfp4_unaligned_hidden_size[swiglu-False-otype0-wtype0-2-2-288-192-1] - AssertionError: Tensor-likes are not close! FAILED tests/moe/test_trtllm_cutlass_fused_moe.py::test_moe_nvfp4_unaligned_hidden_size[swiglu-False-otype0-wtype0-2-2-288-192-4] - AssertionError: Tensor-likes are not close! FAILED tests/moe/test_trtllm_cutlass_fused_moe.py::test_moe_nvfp4_unaligned_hidden_size[swiglu-False-otype0-wtype0-2-2-160-192-1] - AssertionError: Tensor-likes are not close! FAILED tests/moe/test_trtllm_cutlass_fused_moe.py::test_moe_nvfp4_unaligned_hidden_size[swiglu-False-otype0-wtype0-2-2-160-192-4] - AssertionError: Tensor-likes are not close! FAILED tests/moe/test_trtllm_cutlass_fused_moe.py::test_moe_nvfp4_unaligned_hidden_size[swiglu-False-otype0-wtype0-2-2-320-160-1] - AssertionError: Tensor-likes are not close! FAILED tests/moe/test_trtllm_cutlass_fused_moe.py::test_moe_nvfp4_unaligned_hidden_size[swiglu-False-otype0-wtype0-2-2-320-160-4] - AssertionError: Tensor-likes are not close! ================================================================================== 8 failed, 2 passed in 361.01s (0:06:01) ==================================================================================

i.e., the carefully chosen hidden sizes cause failures with NaNs produced.

After reverting the changes in ab0068c, I get back

$ pytest tests/moe/test_trtllm_cutlass_fused_moe.py::test_moe_nvfp4_unaligned_hidden_size -v --tb=short ============================================================================================ test session starts ============================================================================================ ... ... tests/moe/test_trtllm_cutlass_fused_moe.py::test_moe_nvfp4_unaligned_hidden_size[swiglu-False-otype0-wtype0-2-2-288-128-1] PASSED [ 10%] tests/moe/test_trtllm_cutlass_fused_moe.py::test_moe_nvfp4_unaligned_hidden_size[swiglu-False-otype0-wtype0-2-2-288-128-4] PASSED [ 20%] tests/moe/test_trtllm_cutlass_fused_moe.py::test_moe_nvfp4_unaligned_hidden_size[swiglu-False-otype0-wtype0-2-2-288-192-1] PASSED [ 30%] tests/moe/test_trtllm_cutlass_fused_moe.py::test_moe_nvfp4_unaligned_hidden_size[swiglu-False-otype0-wtype0-2-2-288-192-4] PASSED [ 40%] tests/moe/test_trtllm_cutlass_fused_moe.py::test_moe_nvfp4_unaligned_hidden_size[swiglu-False-otype0-wtype0-2-2-160-192-1] PASSED [ 50%] tests/moe/test_trtllm_cutlass_fused_moe.py::test_moe_nvfp4_unaligned_hidden_size[swiglu-False-otype0-wtype0-2-2-160-192-4] PASSED [ 60%] tests/moe/test_trtllm_cutlass_fused_moe.py::test_moe_nvfp4_unaligned_hidden_size[swiglu-False-otype0-wtype0-2-2-320-160-1] PASSED [ 70%] tests/moe/test_trtllm_cutlass_fused_moe.py::test_moe_nvfp4_unaligned_hidden_size[swiglu-False-otype0-wtype0-2-2-320-160-4] PASSED [ 80%] tests/moe/test_trtllm_cutlass_fused_moe.py::test_moe_nvfp4_unaligned_hidden_size[swiglu-False-otype0-wtype0-2-2-256-192-1] PASSED [ 90%] tests/moe/test_trtllm_cutlass_fused_moe.py::test_moe_nvfp4_unaligned_hidden_size[swiglu-False-otype0-wtype0-2-2-256-192-4] PASSED [100%]

which is the current state

bkryu · 2026-04-10T04:17:10Z

/bot run

flashinfer-bot · 2026-04-10T04:17:30Z

GitLab MR !524 has been updated with latest changes, and the CI pipeline #48179740 is currently running. I'll report back once the pipeline job completes.

Ports optimizations from flashinfer-ai/flashinfer#3014 for small-batch decode. For NVFP4/MXFP8 MoE decode with 128 experts and top_k=8, ~120 experts are empty, and the helper kernels were walking all of them or zero-padding scale factors that CUTLASS grouped GEMM never reads. - computeStridesTmaWarpSpecializedKernel: early-exit threads whose assigned expert has gemm_m == 0 after problem-shape init, skipping alpha-scale / block-scaling-factor / stride / pointer setup. - Reduce stride-kernel threadblock from min(1024, experts) to min(32, experts) so 128 experts span 4 SMs instead of 1. - Remove N-dim SF padding loops from expandInputRowsKernel and doActivationKernel and drop the corresponding num_padding_tokens term from both launchers. CUTLASS sets gemm_m = tokens_to_expert per expert and never reads SFs for rows beyond that; K-dim SF padding is still required and remains. Reported ~5-12% decode speedup on RTX Pro 6000 and 2-6% on Spark, neutral elsewhere. Exercised by existing 128-expert NVFP4 / W4A8_NVFP4_FP8 / W4A8_MXFP4_MXFP8 coverage in tests/unittest/_torch/modules/moe/test_moe_backend.py and cpp/tests/unit_tests/kernels/mixtureOfExpertsTest.cu. Signed-off-by: Pamela <179191831+pamelap-nvidia@users.noreply.github.com> [None][test] Cover inactive-expert paths in CUTLASS MoE gtest Adds four TYPED_TESTs to mixtureOfExpertsTest.cu exercising the code paths touched by the prior commit. All four leave mIsLongTest at its default (false), so BasicPermuteTest / ParallelismTest run the memset-poison determinism check (0xD5 vs 0x2A) that validates output independence from uninitialised workspace. - PermuteManyInactiveExperts: k=8, 128 experts, 3 tokens. Most experts receive no tokens, exercising the gemm_m == 0 early-exit and, for block-scaled dtypes (NVFP4, MXFP8xMXFP4), the removed N-dim SF padding region. - PermuteSingleTokenDecode: k=8, 128 experts, 1 token. Every active expert gets gemm_m == 1 (tightest valid case) alongside ~120 inactive experts. - PermuteSwigluManyInactiveExperts: same many-inactive profile with Swiglu activation, exercising the block-scaled GLU path in doActivationKernel on MXFP8xMXFP4. - ExpertParallelManyInactiveExperts: shards 128 experts across 64 or 128 ranks with k=8 and 3 tokens; many ranks see zero active tokens, end-to-end including the alltoall-adjacent flow. Signed-off-by: Pamela <179191831+pamelap-nvidia@users.noreply.github.com> [None][fix] Guard inactive-expert early-exit and Swiglu test Two follow-up fixes for the H100 regression surfaced in CUTLASS MoE gtest: - Gate the gemm_m == 0 early exit in computeStridesTmaWarpSpecializedKernel on !int4_groupwise_params.enabled. The W4A8_AWQ / WFP4A16 CUTLASS kernel instantiations read per-expert stride_s_a / ptr_s_a during init regardless of per-group tile count, so skipping setup left those arrays uninitialised and crashed PermuteSweepNumTokensGeglu / misresulted PermuteSweepNumTokens and PermuteMixtral8x7b on FP8+uint4+bf16. The optimisation still applies to NVFP4, MXFP8xMXFP4, FP8, BF16, and FP16 paths where FlashInfer validated the safety of the early exit. - Skip PermuteSwigluManyInactiveExperts for half output. FP16's narrow dynamic range can't hold top_k=8 Swiglu accumulation at the minimum- alignment hidden size within compareFinal's tolerance. Mirrors the existing NVFP4 Relu-only skip in BasicPermuteTest. BF16 and MXFP8xMXFP4 instances continue to run and retain coverage. Signed-off-by: Pamela <179191831+pamelap-nvidia@users.noreply.github.com> [None][chore] Bump copyright year on modified CUTLASS MoE files Signed-off-by: Pamela <179191831+pamelap-nvidia@users.noreply.github.com>

Ports optimizations from flashinfer-ai/flashinfer#3014 for small-batch decode. For NVFP4/MXFP8 MoE decode with 128 experts and top_k=8, ~120 experts are empty, and the helper kernels were walking all of them or zero-padding scale factors that CUTLASS grouped GEMM never reads. - computeStridesTmaWarpSpecializedKernel: early-exit threads whose assigned expert has gemm_m == 0 after problem-shape init, skipping alpha-scale / block-scaling-factor / stride / pointer setup. - Reduce stride-kernel threadblock from min(1024, experts) to min(32, experts) so 128 experts span 4 SMs instead of 1. - Remove N-dim SF padding loops from expandInputRowsKernel and doActivationKernel and drop the corresponding num_padding_tokens term from both launchers. CUTLASS sets gemm_m = tokens_to_expert per expert and never reads SFs for rows beyond that; K-dim SF padding is still required and remains. Reported ~5-12% decode speedup on RTX Pro 6000 and 2-6% on Spark, neutral elsewhere. Exercised by existing 128-expert NVFP4 / W4A8_NVFP4_FP8 / W4A8_MXFP4_MXFP8 coverage in tests/unittest/_torch/modules/moe/test_moe_backend.py and cpp/tests/unit_tests/kernels/mixtureOfExpertsTest.cu. Signed-off-by: Pamela <179191831+pamelap-nvidia@users.noreply.github.com>

bkryu added 14 commits April 6, 2026 16:42

Add spark to microbenchmarks

0cc6825

CUTLASS MoE

0fb4db3

Tune

a94b434

Fix

2770a20

computeStrides optimize

866dbbc

computeStrides optimize

175b4c8

Skip memset

1b3e421

Cleanup

9bc0e0b

Perf tune

b3d96ab

revert perf tune

2324fb3

Perf Tune v2 -- Flatten N-dim padding

626fed1

Merge branch 'main' into bench_spark

f7eb890

Merge branch 'main' into bench_spark

527ede3

Port optimizations from TRTLLM

3ab2921

bkryu requested review from IwakuraRein, aleozlx, jiahanc, nv-yunzheq, samuellees and yzh119 as code owners April 8, 2026 04:49

flashinfer-bot added the op: moe label Apr 8, 2026

bkryu added the run-ci label Apr 8, 2026

Merge branch 'main' into bench_spark

e9224bf

bkryu changed the title ~~perf: Optimize MoE helper kernels for small-batch decode workloads~~ perf: Optimize CUTLASS MoE helper kernels for small-batch decode workloads Apr 8, 2026

gemini-code-assist Bot reviewed Apr 8, 2026

View reviewed changes

bkryu added 2 commits April 8, 2026 04:54

pre-commit

b5d140b

Address comments

451befd

coderabbitai Bot reviewed Apr 8, 2026

View reviewed changes

Comment thread csrc/fused_moe/cutlass_backend/cutlass_fused_moe_kernels.cuh Outdated

Reduce unnecessary diffs from comments

a33b48c

bkryu added 2 commits April 8, 2026 21:49

NAN fill testing for removing the N-padding

85ceb8a

remove the 0xFF poison memsets

37dc223

coderabbitai Bot reviewed Apr 9, 2026

View reviewed changes

bkryu added 3 commits April 9, 2026 05:02

Merge branch 'main' into bench_spark

45f97f0

Poison with 0xff AFTER kernel to make sure poisoning works

6436f8e

Undo poisoning

a22ce59

yzh119 reviewed Apr 9, 2026

View reviewed changes

bkryu added 2 commits April 9, 2026 12:26

Address comments -- replace PDL PTX instruction to CUDA API

e7bc590

Merge branch 'main' into bench_spark

cabf9cd

nv-yunzheq approved these changes Apr 9, 2026

View reviewed changes

djns99 approved these changes Apr 10, 2026

View reviewed changes

bkryu added 2 commits April 9, 2026 21:03

Poison K-dim for validation

6f36003

Remove the poisoned K-dim padding

ab0068c

aleozlx approved these changes Apr 10, 2026

View reviewed changes

aleozlx merged commit 6248936 into flashinfer-ai:main Apr 10, 2026
33 checks passed

pamelap-nvidia mentioned this pull request Apr 20, 2026

[TRTLLM-10362][perf] Skip inactive experts in CUTLASS MoE stride/SF helpers NVIDIA/TensorRT-LLM#13241

Open

1 task

coderabbitai Bot mentioned this pull request Apr 21, 2026

feat(moe-a2a): Update nvlink onesided all-to-all #3139

Closed

5 tasks

coderabbitai Bot mentioned this pull request May 11, 2026

Support 4over6 nvfp4 for quantizer and fused MoE #3264

Merged

5 tasks

coderabbitai Bot mentioned this pull request May 18, 2026

feat: Add Hopper MXFP4 x FP8 fused MoE support #3349

Open

5 tasks

bkryu deleted the bench_spark branch May 29, 2026 20:43

coderabbitai Bot mentioned this pull request Jun 2, 2026

feat(moe): add SWIGLUSTEP activation to CUTLASS fused MoE #3492

Merged

5 tasks

coderabbitai Bot mentioned this pull request Jul 8, 2026

perf: optimize MXFP4xBF16 & INT4xFP8 and add MXFP4xFP8 CUTLASS MoE backend for SM90 #3738

Merged

8 tasks

	round_up = lambda x, y: (x + y - 1) // y * y
	def round_up(x: int, y: int) -> int:
	return ceil_div(x, y) * y


		flash_output = torch.zeros_like(x)

		# Poison GPU memory with 0xFF — same approach as the NVFP4 test.

Uh oh!

Conversation

bkryu commented Apr 8, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

📌 Description

Summary

Motivation

Changes

Correctness

Large-batch regression safety

Performance Numbers

🔍 Related Issues

🚀 Pull Request Checklist

✅ Pre-commit Checks

🧪 Tests

Reviewer Notes

Summary by CodeRabbit

Uh oh!

coderabbitai Bot commented Apr 8, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Reviews paused

Walkthrough

Changes

Sequence Diagram(s)

Estimated code review effort

Possibly related PRs

Suggested reviewers

Poem

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist Bot Apr 8, 2026

Choose a reason for hiding this comment

Uh oh!

bkryu Apr 8, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Apr 8, 2026

Choose a reason for hiding this comment

Uh oh!

bkryu Apr 8, 2026

Choose a reason for hiding this comment

Uh oh!

coderabbitai Bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

coderabbitai Bot left a comment

Choose a reason for hiding this comment

Uh oh!

coderabbitai Bot Apr 9, 2026

Choose a reason for hiding this comment

Uh oh!

coderabbitai Bot Apr 9, 2026

Choose a reason for hiding this comment

Uh oh!

coderabbitai Bot Apr 9, 2026

Choose a reason for hiding this comment

Uh oh!

bkryu commented Apr 9, 2026

Uh oh!

flashinfer-bot commented Apr 9, 2026

Uh oh!

yzh119 Apr 9, 2026

Choose a reason for hiding this comment

Uh oh!

bkryu Apr 9, 2026

Choose a reason for hiding this comment

Uh oh!

bkryu commented Apr 9, 2026

Uh oh!

flashinfer-bot commented Apr 9, 2026

Uh oh!

djns99 Apr 10, 2026

Choose a reason for hiding this comment

Uh oh!

bkryu Apr 10, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

bkryu commented Apr 8, 2026 •

edited

Loading

coderabbitai Bot commented Apr 8, 2026 •

edited

Loading

bkryu Apr 10, 2026 •

edited

Loading